# 文件名: modules/snake.py
# Keep this file mostly the same, ensure head color calculation is robust
import pygame
import copy
from .utils import draw_grid_cell, grid_to_pixel

'''贪吃蛇类'''
class Snake:
    def __init__(self, cfg, initial_body=None, color=None):
        self.cfg = cfg
        self.color = color if color else cfg.GREEN
        # Calculate a brighter head color safely
        r, g, b = self.color
        self.head_color = (min(r + 70, 255), min(g + 70, 255), min(b + 70, 255))

        if initial_body:
            self.body = copy.deepcopy(initial_body) # Ensure deep copy if provided
        else:
            start_pos = copy.deepcopy(cfg.INITIAL_SNAKE_GRID_POS)
            self.body = [start_pos]
            # Create initial body extending left from the head
            for i in range(1, cfg.INITIAL_SNAKE_LENGTH):
                # Check bounds before adding segment
                new_x = start_pos[0] - i
                if new_x >= 0: # Ensure initial body is within grid
                     self.body.append([new_x, start_pos[1]])
                else: # Should not happen with default pos, but safeguard
                     print("Warning: Initial snake body segment out of bounds.")
                     break # Stop adding segments if out of bounds

        self.direction = "RIGHT" # Initial direction
        self.change_to = self.direction
        self.initial_length = cfg.INITIAL_SNAKE_LENGTH # Store initial length for shrink logic


    def change_direction(self, new_direction):
        """请求改变方向, 防止 180 度转弯"""
        if new_direction == "UP" and self.direction != "DOWN":
            self.change_to = "UP"
        elif new_direction == "DOWN" and self.direction != "UP":
            self.change_to = "DOWN"
        elif new_direction == "LEFT" and self.direction != "RIGHT":
            self.change_to = "LEFT"
        elif new_direction == "RIGHT" and self.direction != "LEFT":
            self.change_to = "RIGHT"

    def update_direction(self):
        """确认并更新方向"""
        # This check prevents the snake from reversing onto itself in the same frame
        # if multiple key presses happen quickly. The check in change_direction
        # handles the logic, this just applies the change.
        self.direction = self.change_to

    def move(self, grow=False, shrink=False):
        """移动蛇，处理增长和缩短"""
        # Get current head position
        head = copy.deepcopy(self.body[0])

        # Calculate new head position based on direction
        if self.direction == "UP":
            head[1] -= 1
        elif self.direction == "DOWN":
            head[1] += 1
        elif self.direction == "LEFT":
            head[0] -= 1
        elif self.direction == "RIGHT":
            head[0] += 1

        # Insert new head
        self.body.insert(0, head)

        # Handle shrinking: Remove tail if shrink=True and length > 1 (don't shrink to nothing)
        if shrink and len(self.body) > 1:
           self.body.pop()
        # Handle normal movement (no growth): Remove tail if grow=False and length > 1
        elif not grow and len(self.body) > 1:
           self.body.pop()
        # If grow=True, do nothing (tail is not removed)
        # If shrink=True but length is already 1, do nothing

    def check_collision(self, traps, mode="wall"):
        """检查碰撞: 墙壁, 自身, 陷阱"""
        if not self.body: return True # No snake body means collision

        head = self.body[0]

        # 1. Check wall collision (if mode includes wall check)
        if mode == "wall":
            if (head[0] < 0 or head[0] >= self.cfg.GRID_WIDTH or
                    head[1] < 0 or head[1] >= self.cfg.GRID_HEIGHT):
                # print(f"Collision: Wall at {head}")
                return True

        # 2. Check self collision (head collides with any part of the body *after* the head)
        if head in self.body[1:]:
            # print(f"Collision: Self at {head}")
            return True

        # 3. Check trap collision
        if head in traps:
            # print(f"Collision: Trap at {head}")
            return True

        return False

    def handle_wrap_around(self):
         """处理穿墙逻辑 (如果游戏模式需要) - Not used by default in this setup"""
         if not self.body: return
         head = self.body[0]
         if head[0] >= self.cfg.GRID_WIDTH: head[0] = 0
         elif head[0] < 0: head[0] = self.cfg.GRID_WIDTH - 1
         if head[1] >= self.cfg.GRID_HEIGHT: head[1] = 0
         elif head[1] < 0: head[1] = self.cfg.GRID_HEIGHT - 1


    def draw(self, screen):
        """绘制蛇，区分蛇头，并添加眼睛"""
        if not self.body: return # Don't draw if snake has no body

        # --- Draw snake body (all segments except the head) ---
        for segment in self.body[1:]:
            draw_grid_cell(screen, segment, self.color, self.cfg)

        # --- Draw snake head ---
        head_pos = self.body[0]
        draw_grid_cell(screen, head_pos, self.head_color, self.cfg)

        # --- Draw eyes on the head ---
        head_pixel_pos = grid_to_pixel(head_pos, self.cfg)
        # Make eye size relative to grid size
        eye_radius = self.cfg.GRID_SIZE // 6
        pupil_radius = eye_radius // 2
        # Offset from corners, adjusted for radius
        offset = self.cfg.GRID_SIZE // 4

        eye_color = self.cfg.WHITE
        pupil_color = self.cfg.BLACK

        # Calculate positions for two eyes based on direction
        if self.direction == "UP":
            eye1_center = (head_pixel_pos[0] + offset, head_pixel_pos[1] + offset)
            eye2_center = (head_pixel_pos[0] + self.cfg.GRID_SIZE - offset, head_pixel_pos[1] + offset)
            pupil_offset = (0, -pupil_radius // 2)
        elif self.direction == "DOWN":
            eye1_center = (head_pixel_pos[0] + offset, head_pixel_pos[1] + self.cfg.GRID_SIZE - offset)
            eye2_center = (head_pixel_pos[0] + self.cfg.GRID_SIZE - offset, head_pixel_pos[1] + self.cfg.GRID_SIZE - offset)
            pupil_offset = (0, pupil_radius // 2)
        elif self.direction == "LEFT":
            eye1_center = (head_pixel_pos[0] + offset, head_pixel_pos[1] + offset)
            eye2_center = (head_pixel_pos[0] + offset, head_pixel_pos[1] + self.cfg.GRID_SIZE - offset)
            pupil_offset = (-pupil_radius // 2, 0)
        else: # RIGHT (default)
            eye1_center = (head_pixel_pos[0] + self.cfg.GRID_SIZE - offset, head_pixel_pos[1] + offset)
            eye2_center = (head_pixel_pos[0] + self.cfg.GRID_SIZE - offset, head_pixel_pos[1] + self.cfg.GRID_SIZE - offset)
            pupil_offset = (pupil_radius // 2, 0)

        # Draw white part of eyes (circles)
        pygame.draw.circle(screen, eye_color, eye1_center, eye_radius)
        pygame.draw.circle(screen, eye_color, eye2_center, eye_radius)

        # Draw pupils (smaller circles, offset)
        pupil1_center = (eye1_center[0] + pupil_offset[0], eye1_center[1] + pupil_offset[1])
        pupil2_center = (eye2_center[0] + pupil_offset[0], eye2_center[1] + pupil_offset[1])
        pygame.draw.circle(screen, pupil_color, pupil1_center, pupil_radius)
        pygame.draw.circle(screen, pupil_color, pupil2_center, pupil_radius)


    @property
    def head_coord(self):
        """Return the coordinates of the snake's head."""
        return self.body[0] if self.body else None

    @property
    def coords(self):
        """Return a list of all coordinates occupied by the snake's body."""
        return self.body